// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <gtest/gtest.h>
#include <pollux/buffer/buffer.h>
#include <pollux/vector/complex_vector.h>
#include <pollux/testing/vector/vector_test_base.h>

using namespace kumo::pollux;
using namespace kumo::pollux::test;

class IsWritableVectorTest : public testing::Test, public VectorBuilder {
 protected:
  static void SetUpTestCase() {
    memory::MemoryManager::testingSetInstance({});
  }

  // We use templates here to avoid the compiler automatically creating new
  // shared_ptrs which it would do if we used VectorPtr.
  template <typename T, typename V>
  void testChildVector(
      const std::shared_ptr<T>& vector,
      const std::shared_ptr<V>& childVector) {
    ASSERT_TRUE(vector->is_nulls_writable());
    ASSERT_TRUE(BaseVector::is_vector_writable(vector));

    {
      // Make a copy of childVector so it's no longer writable.
      auto copy = childVector;
      ASSERT_TRUE(vector->is_nulls_writable());
      ASSERT_FALSE(BaseVector::is_vector_writable(vector));
    }
  }

  // We use templates here to avoid the compiler automatically creating new
  // shared_ptrs which it would do if we used VectorPtr.
  template <typename T>
  void testBufferPtr(
      const std::shared_ptr<T>& vector,
      const BufferPtr& buffer) {
    ASSERT_TRUE(vector->is_nulls_writable());
    ASSERT_TRUE(BaseVector::is_vector_writable(vector));

    {
      // Make a copy of buffer so it's no longer mutable.
      auto copy = buffer;
      ASSERT_TRUE(vector->is_nulls_writable());
      ASSERT_FALSE(BaseVector::is_vector_writable(vector));
    }

    // Make buffer multiply-referenced so it's no longer be mutable.
    auto copy = buffer;
    ASSERT_TRUE(vector->is_nulls_writable());
    ASSERT_FALSE(BaseVector::is_vector_writable(vector));

    copy = nullptr;

    // Make sure nothing gets left unwritable.
    ASSERT_TRUE(vector->is_nulls_writable());
    ASSERT_TRUE(BaseVector::is_vector_writable(vector));
  }

  // We use templates here to avoid the compiler automatically creating new
  // shared_ptrs which it would do if we used VectorPtr.
  template <typename T>
  void basicTest(const std::shared_ptr<T>& vector) {
    ASSERT_TRUE(vector->is_nulls_writable());
    ASSERT_TRUE(BaseVector::is_vector_writable(vector));

    {
      auto copy = vector;
      ASSERT_TRUE(vector->is_nulls_writable());
      // The Vector is not uniquely referenced.
      ASSERT_FALSE(BaseVector::is_vector_writable(vector));
    }

    // Hack to ensure nulls are allocated.
    vector->mutable_raw_nulls();
    ASSERT_TRUE(vector->is_nulls_writable());
    ASSERT_TRUE(BaseVector::is_vector_writable(vector));

    {
      // Make a copy of nulls so it's no longer mutable.
      auto nulls = vector->nulls();
      ASSERT_FALSE(vector->is_nulls_writable());
      ASSERT_FALSE(BaseVector::is_vector_writable(vector));
    }
  }
};

TEST_F(IsWritableVectorTest, flat_vector) {
  auto flat_vector = make_flat_vector<int32_t>({1, 2, 3});

  basicTest(flat_vector);
  testBufferPtr(flat_vector, flat_vector->values());
}

TEST_F(IsWritableVectorTest, array_vector) {
  auto array_vector = make_array_vector<int32_t>({{1, 2}, {3, 4}, {5, 6}});

  basicTest(array_vector);
  testBufferPtr(array_vector, array_vector->offsets());
  testBufferPtr(array_vector, array_vector->sizes());
  testChildVector(array_vector, array_vector->elements());
}

TEST_F(IsWritableVectorTest, map_vector) {
  auto map_vector = make_map_vector(
      {0, 2, 4, 6},
      make_flat_vector<int32_t>({1, 2, 3, 4, 5, 6}),
      make_flat_vector<int32_t>({1, 2, 3, 4, 5, 6}));

  basicTest(map_vector);
  testBufferPtr(map_vector, map_vector->offsets());
  testBufferPtr(map_vector, map_vector->sizes());
  testChildVector(map_vector, map_vector->mapKeys());
  testChildVector(map_vector, map_vector->mapValues());
}

TEST_F(IsWritableVectorTest, row_vector) {
  auto row_vector = make_row_vector(
      {make_flat_vector<int32_t>({1, 2, 3, 4, 5, 6}),
       make_flat_vector<int32_t>({1, 2, 3, 4, 5, 6}),
       make_flat_vector<int32_t>({1, 2, 3, 4, 5, 6})});

  basicTest(row_vector);
  testChildVector(row_vector, row_vector->childAt(0));
  testChildVector(row_vector, row_vector->childAt(1));
  testChildVector(row_vector, row_vector->childAt(2));
}
